"""构建Seq2seq模型, 将编码器和解码器输出的结果进行合并"""
import torch.nn as nn

from decoder import Decoder
from encoder import Encoder
import config


class Seq2seq(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()

    def forward(self, input, input_length):
        # _1: [batch_size, max_len, hidden_size]
        # _2: 每个数字转换成字符串或列表后原始的长度
        # encoder_hidden: [1, batch_size, hidden_size]
        _1, _2, encoder_hidden = self.encoder(input, input_length)
        # decoder_output: [batch_size, max_len+2, vocab_size]
        # decoder_hidden: [1, batch_size, hidden_size]
        decoder_output, decoder_hidden = self.decoder(encoder_hidden)

        return decoder_output

    # def evaluate(self, input, input_length):
    #     _1, _2, encoder_hidden = self.encoder(input, input_length)
    #     decoder_predict = self.decoder.evaluate(encoder_hidden)
    #
    #     return decoder_predict

    def evaluate(self, input, input_length):
        _1, _2, encoder_hidden = self.encoder(input, input_length)
        decoder_predict = self.decoder.evaluate(encoder_hidden)

        # decoder_predict: [max_len+2, batch_size]
        return decoder_predict





